from flask import Flask, request, jsonify, render_template
import numpy as np
import joblib
import os
from sklearn.neural_network import MLPClassifier
from sklearn.preprocessing import StandardScaler
import json

app = Flask(__name__)

# 全局变量存储模型
model = None
scaler = None
training_data = []
training_labels = []

def init_model():
    """初始化或加载模型"""
    global model, scaler
    
    # 检查是否存在已保存的模型
    if os.path.exists('static/models/digit_classifier.pkl'):
        model_data = joblib.load('static/models/digit_classifier.pkl')
        model = model_data['model']
        scaler = model_data['scaler']
        print("已加载现有模型")
    else:
        # 创建新模型
        model = MLPClassifier(
            hidden_layer_sizes=(128, 64),
            activation='relu',
            solver='adam',
            max_iter=1000,
            random_state=42
        )
        scaler = StandardScaler()
        print("已创建新模型")
    
    return model, scaler

@app.route('/')
def index():
    """渲染主页"""
    return render_template('index.html')

@app.route('/api/predict', methods=['POST'])
def predict():
    """预测手写数字"""
    try:
        data = request.json
        pixels = np.array(data['pixels']).reshape(1, -1)
        
        if model is None or scaler is None:
            init_model()
        
        # 标准化输入
        pixels_scaled = scaler.transform(pixels)
        
        # 预测
        prediction = model.predict(pixels_scaled)[0]
        probabilities = model.predict_proba(pixels_scaled)[0]
        confidence = float(max(probabilities))
        
        return jsonify({
            'prediction': int(prediction),
            'confidence': confidence,
            'probabilities': probabilities.tolist()
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/train', methods=['POST'])
def train():
    """训练模型"""
    try:
        data = request.json
        pixels = np.array(data['pixels']).reshape(1, -1)
        label = int(data['label'])
        
        # 添加到训练数据
        training_data.append(pixels[0])
        training_labels.append(label)
        
        # 如果数据足够，开始训练
        if len(training_data) >= 5:
            X = np.array(training_data)
            y = np.array(training_labels)
            
            # 标准化
            X_scaled = scaler.fit_transform(X)
            
            # 训练模型
            model.fit(X_scaled, y)
            
            # 计算准确率
            accuracy = model.score(X_scaled, y)
            
            # 保存模型
            model_data = {'model': model, 'scaler': scaler}
            joblib.dump(model_data, 'static/models/digit_classifier.pkl')
            
            return jsonify({
                'status': 'success',
                'trained_samples': len(training_data),
                'accuracy': float(accuracy),
                'message': f'模型已更新，准确率: {accuracy:.2%}'
            })
        else:
            return jsonify({
                'status': 'collecting',
                'trained_samples': len(training_data),
                'message': f'已收集 {len(training_data)} 个样本，需要至少5个样本才能训练'
            })
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/reset', methods=['POST'])
def reset():
    """重置模型"""
    global training_data, training_labels
    
    training_data = []
    training_labels = []
    
    # 删除保存的模型文件
    if os.path.exists('static/models/digit_classifier.pkl'):
        os.remove('static/models/digit_classifier.pkl')
    
    # 重新初始化模型
    init_model()
    
    return jsonify({'status': 'success', 'message': '模型已重置'})

@app.route('/api/status')
def status():
    """获取模型状态"""
    return jsonify({
        'trained_samples': len(training_data),
        'model_exists': os.path.exists('static/models/digit_classifier.pkl')
    })

if __name__ == '__main__':
    init_model()
    app.run(debug=True, host='0.0.0.0', port=5000)
